from contextlib import asynccontextmanager

from fastapi import FastAPI
from tortoise.contrib.fastapi import RegisterTortoise
from tortoise import Tortoise, run_async, connections

from config import config

MYSQL_HOST = config.MYSQL_HOST
MYSQL_PORT = config.MYSQL_PORT
MYSQL_USER = config.MYSQL_USER
MYSQL_PASSWORD = config.MYSQL_PASSWORD
MYSQL_DATABASE = config.MYSQL_DATABASE

TORTOISE_MODELS = ['account.models']


TORTOISE_CONFIG = {
    'connections': {
        # 'default': {
        #     'engine': 'tortoise.backends.mysql',
        #     'credentials': {
        #         'host': "127.0.0.1",
        #         'port': MYSQL_PORT,
        #         'user': MYSQL_USER,
        #         'password': MYSQL_PASSWORD,
        #         'database': MYSQL_DATABASE,
        #     }
        # },
        'LA Connection': {
            'engine': 'tortoise.backends.mysql',
            'credentials': {
                'host': "127.0.0.1",
                'port': MYSQL_PORT,
                'user': MYSQL_USER,
                'password': MYSQL_PASSWORD,
                'database': "la_db",
            }
        }
    },
    'apps': {
        # 'test_models': {
        #     'models': TORTOISE_MODELS,
        #     'default_connection': 'default'
        # },
        'la_models': {
            'models': TORTOISE_MODELS,
            'default_connection': 'LA Connection'
        }
    },
    'use_tz': False,
    'timezone': 'Asia/Shanghai'
}


# @asynccontextmanager
# async def register_mysql(app: FastAPI):
#     try:
#         await Tortoise.init(config=TORTOISE_CONFIG)
#         print("Mysql connected")
#         await Tortoise.generate_schemas(safe=True)
#         print("Tables generated")
#         yield
#         await Tortoise.close_connections()
#         print("Mysql closed")
#     except Exception as e:
#         print(f"Mysql error: {e}")

@asynccontextmanager
async def register_mysql(app: FastAPI):
    try:
        async with RegisterTortoise(
            app,
            config=TORTOISE_CONFIG,
            generate_schemas=True,
            add_exception_handlers=True,
        ):
            print("Mysql connected")
            yield
            await connections.close_all()
            print("Mysql closed")

    except Exception as e:
        print(f"Init database error {e}")
